import subprocess
import tkinter as tk
from tkinter import messagebox, ttk
import logging
import ctypes

class USBIPDManagerApp:
    def __init__(self, master):
        self.master = master
        self.master.title("USB2WSL TOOL")
        self.master.geometry("400x400")  # 增加窗口高度以适应新元素
        self.selected_device = None
        self.wsl_distro = "Ubuntu-22.04"  # 默认值

        # 设置日志记录为UTF-8编码
        logging.basicConfig(
            filename="usbipd_manager.log",
            level=logging.DEBUG,
            format="%(asctime)s - %(levelname)s - %(message)s",
            encoding='utf-8'  # 指定编码格式为UTF-8
        )
        logging.info("启动 USBIPD 管理器")

        # 检查是否以管理员权限运行
        if not self.is_admin():
            messagebox.showwarning("警告", "请以管理员权限运行该程序。")
            logging.warning("程序未以管理员权限运行")
            self.master.quit()
            return

        # 检查usbipd是否安装
        if not self.check_usbipd_installed():
            messagebox.showerror("错误", "未安装usbipd或未在PATH中找到它。请安装usbipd-win。")
            logging.error("usbipd未安装或在PATH中未找到")
            self.master.quit()
            return

        # 主框架
        self.main_frame = ttk.Frame(master, padding="10")
        self.main_frame.grid(sticky=(tk.W, tk.E, tk.N, tk.S))

        # 设置行和列的权重以便于自适应窗口大小
        self.main_frame.grid_rowconfigure(0, weight=1)
        self.main_frame.grid_rowconfigure(1, weight=1)
        self.main_frame.grid_rowconfigure(2, weight=1)
        self.main_frame.grid_rowconfigure(3, weight=1)  # 新增行
        self.main_frame.grid_columnconfigure(0, weight=1)

        # 标题标签
        self.title_label = ttk.Label(self.main_frame, text="USB2WSL TOOL", font=("Arial", 16))
        self.title_label.grid(row=0, column=0, pady=5)

        # 描述标签
        self.description_label = ttk.Label(self.main_frame, text="选择要共享的USB设备", font=("Arial", 10))
        self.description_label.grid(row=2, column=0, pady=5)

        # WSL发行版名称输入标签
        self.wsl_label = ttk.Label(self.main_frame, text="输入WSL发行版名称:", font=("Arial", 10))
        self.wsl_label.grid(row=1, column=0, pady=5)

        # WSL发行版名称输入框
        self.wsl_entry = ttk.Entry(self.main_frame, width=30)
        self.wsl_entry.grid(row=3, column=0, pady=5)
        self.wsl_entry.insert(0, "Ubuntu-22.04")  # 设置默认值

        # USB设备列表框
        self.device_listbox = tk.Listbox(self.main_frame, width=50)
        self.device_listbox.grid(row=4, column=0, pady=10, sticky=(tk.W, tk.E))

        # 按钮框架
        self.button_frame = ttk.Frame(self.main_frame)
        self.button_frame.grid(row=5, column=0, pady=10)

        # 刷新按钮
        self.refresh_button = ttk.Button(self.button_frame, text="刷新USB设备", command=self.refresh_usb_list)
        self.refresh_button.pack(side=tk.LEFT, padx=5)

        # 共享按钮
        self.attach_button = ttk.Button(self.button_frame, text="共享选中设备", command=self.attach_selected_device)
        self.attach_button.pack(side=tk.LEFT, padx=5)

        # 取消共享按钮
        self.detach_button = ttk.Button(self.button_frame, text="取消共享设备", command=self.detach_selected_device)
        self.detach_button.pack(side=tk.LEFT, padx=5)

        # 初始化刷新USB设备列表
        self.refresh_usb_list()

    def is_admin(self):
        """检查是否以管理员权限运行。"""
        try:
            return ctypes.windll.shell32.IsUserAnAdmin()
        except:
            return False

    def check_usbipd_installed(self):
        """检查usbipd是否已安装。"""
        try:
            subprocess.check_output("usbipd --version", shell=True)
            return True
        except subprocess.CalledProcessError:
            return False
        except FileNotFoundError:
            return False

    def refresh_usb_list(self):
        """刷新USB设备列表。"""
        self.device_listbox.delete(0, tk.END)
        try:
            command = "usbipd list"
            logging.debug(f"执行命令: {command}")
            output = subprocess.check_output(command, shell=True).decode("utf-8")
            devices = output.splitlines()
            for device in devices:
                self.device_listbox.insert(tk.END, device)
            logging.info("USB设备成功刷新")
        except subprocess.CalledProcessError as e:
            messagebox.showerror("错误", "无法列出USB设备。请确保usbipd已安装并运行。")
            logging.error(f"无法列出USB设备: {e}")

    def attach_selected_device(self):
        """共享选中的USB设备。"""
        selection = self.device_listbox.curselection()
        if not selection:
            messagebox.showwarning("警告", "未选择设备。")
            logging.warning("未选择用于共享的设备")
            return

        device_info = self.device_listbox.get(selection[0])
        busid = device_info.split()[0]

        # 更新self.wsl_distro的值
        self.wsl_distro = self.wsl_entry.get()

        # 检查设备是否已共享
        if self.is_device_shared(busid):
            messagebox.showinfo("信息", f"设备 {busid} 已处于共享状态。")
            logging.info(f"设备 {busid} 已处于共享状态，跳过共享操作。")
            return

        # 检查设备是否已绑定
        if not self.is_device_bound(busid):
            self.bind_device(busid)

        self.share_device(busid)

    def detach_selected_device(self):
        """取消共享选中的USB设备。"""
        selection = self.device_listbox.curselection()
        if not selection:
            messagebox.showwarning("警告", "未选择设备。")
            logging.warning("未选择用于取消共享的设备")
            return

        device_info = self.device_listbox.get(selection[0])
        busid = device_info.split()[0]

        # 检查设备是否处于attached状态，取消attached和绑定
        if self.is_device_shared(busid):
            self.detach_device(busid)

        # 检查设备是否处于bound状态，如果是，则取消绑定
        if self.is_device_bound(busid) and not self.is_device_shared(busid):
            self.unbind_device(busid)

    def is_device_bound(self, busid):
        """检查设备是否已绑定。"""
        try:
            command = "usbipd list"
            output = subprocess.check_output(command, shell=True).decode("utf-8")
            devices = output.splitlines()
            for device in devices:
                if busid in device and "bound" in device:  # 确保设备处于绑定状态
                    return True
            return False
        except subprocess.CalledProcessError as e:
            logging.error(f"检查设备绑定状态失败: {e}")
            return False

    def bind_device(self, busid):
        """绑定设备以便于共享。"""
        command = f"usbipd bind --busid {busid}"
        try:
            logging.debug(f"执行命令: {command}")
            subprocess.check_call(command, shell=True)
            messagebox.showinfo("成功", f"设备 {busid} 成功绑定。")
            logging.info(f"设备 {busid} 成功绑定")
        except subprocess.CalledProcessError as e:
            messagebox.showerror("错误", f"绑定设备 {busid} 失败。")
            logging.error(f"绑定设备 {busid} 失败: {e}")

    def is_device_shared(self, busid):
        """检查设备是否已共享。"""
        try:
            command = "usbipd list"
            output = subprocess.check_output(command, shell=True).decode("utf-8")
            devices = output.splitlines()
            for device in devices:
                if busid in device and "Attached" in device:  # 确保设备处于共享状态
                    return True
            return False
        except subprocess.CalledProcessError as e:
            logging.error(f"检查设备共享状态失败: {e}")
            return False

    def share_device(self, busid):
        """共享设备。"""
        share_command = f"usbipd attach --busid {busid} --wsl {self.wsl_distro}"
        try:
            logging.debug(f"执行命令: {share_command}")
            subprocess.check_call(share_command, shell=True)
            messagebox.showinfo("成功", f"设备 {busid} 成功共享到 WSL 发行版 {self.wsl_distro}。")
            logging.info(f"设备 {busid} 成功共享到 WSL 发行版 {self.wsl_distro}")
        except subprocess.CalledProcessError as e:
            messagebox.showerror("错误", f"共享设备 {busid} 失败。")
            logging.error(f"共享设备 {busid} 失败: {e}")

    def detach_device(self, busid):
        """取消共享设备。"""
        detach_command = f"usbipd detach --busid {busid}"
        try:
            logging.debug(f"执行命令: {detach_command}")
            subprocess.check_call(detach_command, shell=True)
            messagebox.showinfo("成功", f"设备 {busid} 成功取消共享。")
            logging.info(f"设备 {busid} 成功取消共享")
        except subprocess.CalledProcessError as e:
            messagebox.showerror("错误", f"取消共享设备 {busid} 失败。")
            logging.error(f"取消共享设备 {busid} 失败: {e}")

# 启动应用
if __name__ == "__main__":
    root = tk.Tk()
    app = USBIPDManagerApp(root)
    root.mainloop()
